from flask import Flask, Response, jsonify, request
from config import config
import json
from index_manager import indexManager
from werkzeug.utils import secure_filename

app = Flask(__name__)
from utils import is_supported_file, format_node, format_response
import os


@app.route("/")
def home():
    data = {"version": "1.0.0", "message": "llamaindex智能文档问答系统"}
    return Response(json.dumps(data, ensure_ascii=False), mimetype="application/json")


@app.route("/uploadFile", methods=["POST"])
def uploadFile():
    if "file" not in request.files:
        return jsonify({"error": "请提供文件"}, 400)
    uploaded_file = request.files["file"]
    if uploaded_file.filename == "":
        return jsonify({"error": "未选择文件"}, 400)
    if not is_supported_file(uploaded_file.filename):
        return jsonify({"error": "不支持的文件格式"}, 400)
    try:
        filename = secure_filename(uploaded_file.filename)
        filepath = os.path.join(config.DOCUMENTS_DIR, filename)
        uploaded_file.save(filepath)
        indexManager.insert_document(filepath, filename)
        return jsonify({"success": True, "message": f"文件上传成功"})
    except Exception as e:
        return jsonify({"error": f"上传失败:{str(e)}"}, 500)


@app.route("/getDocuments", methods=["GET"])
def getDocuments():
    try:
        documents = indexManager.get_document_list()
        return jsonify(documents)
    except Exception as e:
        return jsonify({"error": f"获取文件列表失败:{str(e)}"}, 500)


@app.route("/query", methods=["GET"])
def query():
    query_text = request.args.get("question")  # ?question=1+1=?
    if not query_text:
        return jsonify({"error": "请提供查询文本"}, 400)
    try:
        response = indexManager.query_index(query_text)
        sources = [format_node(node) for node in response.source_nodes]
        return jsonify(format_response(str(response), sources))
        # return jsonify({"text": text, "sources": sources})
    except Exception as e:
        return jsonify({"error": f"查询失败:{str(e)}"}, 500)


def main():
    print("Hello from llamaindex!")


if __name__ == "__main__":
    os.makedirs(config.DOCUMENTS_DIR, exist_ok=True)
    app.run(host=config.FLASH_HOST, port=config.FLASH_PORT, debug=True)
